﻿#define CREATE_OPTIX_DATA

/*
* Copyright (c) 2012 Kyle Hayward
* 
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
* 
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
* 
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
* THE SOFTWARE.
*/

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.InteropServices;

using OptixDotNet;
using OptixDotNet.Math;
using OptixDotNetUtils;

//using SlimDX;
using SlimDX.Direct3D11;
using SlimDX.DXGI;
using SlimDX.Windows;
using SlimDX.D3DCompiler;

using Device = SlimDX.Direct3D11.Device;
using Buffer = SlimDX.Direct3D11.Buffer;
using DataStream = SlimDX.DataStream;
using Vector4 = SlimDX.Vector4;

using OptixFormat = OptixDotNet.Format;

namespace InstantRadiosity
{
    struct PosNormal
    {
        public Vector3 Position;
        public Vector3 Normal;

        public static int SizeInBytes { get { return Vector3.SizeInBytes * 2; } }
        
        public PosNormal( Vector3 pos, Vector3 normal )
        {
            Position = pos;
            Normal = normal;
        }
    }

    public class D3DOBJLoader : OBJLoader
    {
        public Shader Shader;

        public string IntersecitonProgPath { get; set; }
        public string IntersecitonProgName { get; set; }
        public string BoundingBoxProgPath { get; set; }
        public string BoundingBoxProgName { get; set; }

        public Context Context { get; set; }
        public Material Material { get; set; }
        public GeometryGroup GeoGroup { get; set; }

        public AccelBuilder Builder { get; set; }
        public AccelTraverser Traverser { get; set; }

        private Renderer mRenderer;

        public D3DOBJLoader( string filename, Renderer renderer )
            : base( filename )
        {
            mRenderer = renderer;

            Builder = AccelBuilder.Bvh;
            Traverser = AccelTraverser.Bvh;
        }

        protected override void CreateGeometry()
        {
            base.CreateGeometry();

            //create global vertex buffer that is shared by all the packets
            //NOTE: only geometry with = number of positions and normals is supported!!!
            DataStream stream = CreateInterleavedVertices();

            Buffer vertexBuffer = new Buffer( mRenderer.Device, stream, new BufferDescription()
            {
                BindFlags = BindFlags.VertexBuffer,
                CpuAccessFlags = CpuAccessFlags.None,
                OptionFlags = ResourceOptionFlags.None,
                SizeInBytes = PosNormal.SizeInBytes * Positions.Count,
                Usage = ResourceUsage.Default
            } );

#if CREATE_OPTIX_DATA
            //create buffer descriptions
            D3DBuffer vBuffer = new D3DBuffer( Context, new D3DBufferDesc()
            {
                Api = D3DApi.Direct3D11,
                Width = (uint)Positions.Count,
                ElemSize = (uint)PosNormal.SizeInBytes,
                Resource = vertexBuffer.ComPointer,
                Type = BufferType.Input,
                Format = OptixDotNet.Format.User,
            } );
#endif

            List<GeometryInstance> instances = new List<GeometryInstance>();
            foreach ( ObjGroup group in Groups )
            {
                //empty group
                if ( group.VIndices.Count == 0 && group.NIndices.Count == 0 && group.TIndices.Count == 0 )
                    continue;

                ValidateGroup( group );

                Packet packet = new Packet();
                packet.VertexBuffer = vertexBuffer;
                packet.VertexCount = Positions.Count;
                packet.VertexOffset = 0;
                packet.VertexStride = PosNormal.SizeInBytes;

                stream.Dispose();
                stream = new DataStream( Int3.SizeInBytes * group.VIndices.Count, true, true );
                Int3[] indices = group.VIndices.ToArray();
                stream.WriteRange<Int3>( indices );
                stream.Position = 0;

                packet.Topology = PrimitiveTopology.TriangleList;
                packet.IndexCount = group.VIndices.Count * 3;
                packet.IndexOffset = 0;
                packet.IndexBuffer = new Buffer( mRenderer.Device, stream, new BufferDescription()
                {
                    BindFlags = BindFlags.IndexBuffer,
                    CpuAccessFlags = CpuAccessFlags.None,
                    OptionFlags = ResourceOptionFlags.None,
                    Usage = ResourceUsage.Default,
                    SizeInBytes = Int3.SizeInBytes * group.VIndices.Count
                } );

#if CREATE_OPTIX_DATA
                //create optix index buffer
                D3DBuffer viBuffer = new D3DBuffer( Context, new D3DBufferDesc() 
                { 
                    Width = (uint)group.VIndices.Count, 
                    Format = OptixFormat.Int3, 
                    Type = BufferType.Input, 
                    Api = D3DApi.Direct3D11,
                    Resource = packet.IndexBuffer.ComPointer 
                } );

                //create a geometry node and set the buffers
                Geometry geometry = new Geometry( Context );
                geometry.IntersectionProgram = new Program( Context, IntersecitonProgPath, IntersecitonProgName );
                geometry.BoundingBoxProgram = new Program( Context, BoundingBoxProgPath, BoundingBoxProgName );
                geometry.PrimitiveCount = (uint)group.VIndices.Count;

                geometry[ "vertex_buffer" ].Set( vBuffer );
                geometry[ "vindex_buffer" ].Set( viBuffer );

                //create a geometry instance
                GeometryInstance instance = new GeometryInstance( Context );
                instance.Geometry = geometry;
                instance.AddMaterial( Material );
#endif

                packet.Material = new D3DMaterial();
                if ( group.mtrl != null )
                {
                    ObjMaterial mtrl = mMtrls[ group.mtrl ];
                    packet.Material.Diffuse = new Vector4( mtrl.Kd.X, mtrl.Kd.Y, mtrl.Kd.Z, 1.0f );
                    packet.Material.Emission = new Vector4( mtrl.Ke.X, mtrl.Ke.Y, mtrl.Ke.Z, 1.0f );

#if CREATE_OPTIX_DATA
                    instance[ "diffuse_color" ].SetFloat3( mtrl.Kd );
                    instance[ "emission_color" ].SetFloat3( mtrl.Ke );
#endif
                }
                else
                {
                    packet.Material.Diffuse = new Vector4( 1, 1, 1, 1 );
                    packet.Material.Emission = Vector4.Zero;

#if CREATE_OPTIX_DATA
                    instance[ "diffuse_color" ].Set( 1.0f, 1.0f, 1.0f );
#endif
                }
#if CREATE_OPTIX_DATA
                instances.Add( instance );
#endif

                int sizeInBytes = Marshal.SizeOf( typeof( D3DMaterial ) );
                packet.MtrlBuffer = new Buffer( mRenderer.Device, new BufferDescription()
                {
                    BindFlags = BindFlags.ConstantBuffer,
                    CpuAccessFlags = CpuAccessFlags.Write,
                    Usage = ResourceUsage.Dynamic,
                    OptionFlags = ResourceOptionFlags.None,
                    SizeInBytes = sizeInBytes,
                    StructureByteStride = 0
                } );

                SlimDX.DataBox data = mRenderer.Device.ImmediateContext.MapSubresource( packet.MtrlBuffer, MapMode.WriteDiscard, SlimDX.Direct3D11.MapFlags.None );
                data.Data.Write<D3DMaterial>( packet.Material );
                mRenderer.Device.ImmediateContext.UnmapSubresource( packet.MtrlBuffer, 0 );

                packet.Shader = Shader;

                mRenderer.AddPacket( ref packet );
            }

#if CREATE_OPTIX_DATA
            //create an acceleration structure for the geometry
            Acceleration accel = new Acceleration( Context, Builder, Traverser );

            if ( Builder == AccelBuilder.Sbvh || Builder == AccelBuilder.TriangleKdTree )
            {
                accel.VertexBufferName = "vertex_buffer";
                accel.IndexBufferName = "vindex_buffer";
            }

            //now attach the instance and accel to the geometry group
            GeoGroup.Acceleration = accel;
            GeoGroup.AddChildren( instances );
#endif
        }

        private DataStream CreateInterleavedVertices()
        {
            DataStream stream = new DataStream( PosNormal.SizeInBytes * Positions.Count, true, true );

            for ( int i = 0; i < Positions.Count; i++ )
            {
                stream.Write<Vector3>( Positions[ i ] );
                stream.Write<Vector3>( Normals[ i ] );
            }
            stream.Position = 0;

            return stream;
        }
    }
}
